import os
import re
from dataclasses import dataclass
from typing import Any, Optional, Tuple

import torch


@dataclass
class LMCacheEngineMetadata:
    """name of the LLM model"""

    model_name: str
    """ world size when running under a distributed setting """
    world_size: int
    """ worker id when running under a distributed setting """
    worker_id: int
    """ the format of kv tensors """
    fmt: str
    """ the data type of kv tensors """
    kv_dtype: torch.dtype
    """ the shape of kv tensors """
    """ (num_layer, 2, chunk_size, num_kv_head, head_size) """
    kv_shape: tuple[int, int, int, int, int]
    """ whether use MLA"""
    use_mla: bool = False


@dataclass
class DiskCacheMetadata:
    path: str
    size: int  # in bytes
    shape: Optional[torch.Size] = None
    dtype: Optional[torch.dtype] = None
    is_pin: bool = False

    def pin(self) -> bool:
        self.is_pin = True
        return True

    def unpin(self) -> bool:
        self.is_pin = False
        return True

    @property
    def is_pinned(self) -> bool:
        return self.is_pin


@dataclass
class LMCacheMemPoolMetadata:
    """Subset of `LMCacheEngineMetadata` to initialize MemPool"""

    kv_shape: Tuple[int, int, int, int, int]
    kv_dtype: torch.dtype
    max_local_cache_size: int


blend_default_separator = "[BLEND_SEP]"


@dataclass
class LMCacheEngineConfig:
    chunk_size: int
    local_cpu: bool
    max_local_cpu_size: float  # in GB
    # need to be assigned a non-zero
    # value even if local_cpu is disabled
    local_disk: Optional[str]
    max_local_disk_size: float  # in GB

    remote_url: Optional[str]
    remote_serde: Optional[str]  # Can be "naive" or "cachegen"

    save_decode_cache: bool  # whether to store decode kv cache

    # Blending related configurations
    enable_blending: bool  # whether to enable blending
    blend_recompute_ratio: float  # the ratio of blending recompute
    blend_min_tokens: int  # the minimum number of tokens for blending
    blend_special_str: str = " # # "  # the separator for blending

    # P2P related configurations
    enable_p2p: bool = False  # whether to enable peer-to-peer sharing
    lookup_url: Optional[str] = None  # the url of the lookup server
    distributed_url: Optional[str] = None  # the url of the distributed server

    # Error handling related configurations
    error_handling: bool = False  # whether to enable error handling

    # Controller related configurations
    enable_controller: Optional[bool] = False  # whether to enable controller
    # the id of the lmcache instance
    lmcache_instance_id: str = "lmcache_default_instance"
    # controller url
    controller_url: Optional[str] = None
    # lmcache worker url
    # NOTE: port number will add `worker_id`
    lmcache_worker_port: Optional[int] = None

    # (Optional) Nixl configurations
    # whether to enable Nixl
    enable_nixl: Optional[bool] = False
    # Role: sender or receiver
    nixl_role: Optional[str] = None
    # The host of the nixl receiver
    nixl_receiver_host: Optional[str] = None
    # The BASE port of the nixl receiver,
    # real port is nixl_receiver_port + WORKER_RANK
    nixl_receiver_port: Optional[int] = None
    # The transport buffer size of nixl in bytes
    nixl_buffer_size: Optional[int] = None
    # The device that nixl uses
    nixl_buffer_device: Optional[str] = None
    # HACK: explicit option to enable/disable nixl GC before it's mature enough
    nixl_enable_gc: Optional[bool] = False

    # The url of the actual remote lmcache instance for auditing
    audit_actual_remote_url: Optional[str] = None

    @staticmethod
    def from_defaults(
        chunk_size: int = 256,
        local_cpu: bool = True,
        max_local_cpu_size: float = 5.0,
        local_disk: Optional[str] = None,
        max_local_disk_size: int = 0,
        remote_url: Optional[str] = "lm://localhost:65432",
        remote_serde: Optional[str] = "naive",
        save_decode_cache: bool = False,
        enable_blending: bool = False,
        blend_recompute_ratio: float = 0.15,
        blend_min_tokens: int = 256,
        blend_special_str: str = " # # ",
        enable_p2p: bool = False,
        lookup_url: Optional[str] = None,
        distributed_url: Optional[str] = None,
        error_handling: bool = False,
        enable_controller: Optional[bool] = False,
        lmcache_instance_id: str = "lmcache_default_instance",
        controller_url: Optional[str] = None,
        lmcache_worker_port: Optional[int] = None,
        enable_nixl: Optional[bool] = False,
        nixl_role: Optional[str] = None,
        nixl_receiver_host: Optional[str] = None,
        nixl_receiver_port: Optional[int] = None,
        nixl_buffer_size: Optional[int] = None,
        nixl_buffer_device: Optional[str] = None,
        nixl_enable_gc: Optional[bool] = False,
        audit_actual_remote_url: Optional[str] = None,
    ) -> "LMCacheEngineConfig":
        # TODO (ApostaC): Add nixl config
        return LMCacheEngineConfig(
            chunk_size,
            local_cpu,
            max_local_cpu_size,
            local_disk,
            max_local_disk_size,
            remote_url,
            remote_serde,
            save_decode_cache,
            enable_blending,
            blend_recompute_ratio,
            blend_min_tokens,
            blend_special_str,
            enable_p2p,
            lookup_url,
            distributed_url,
            error_handling,
            enable_controller,
            lmcache_instance_id,
            controller_url,
            lmcache_worker_port,
            enable_nixl,
            nixl_role,
            nixl_receiver_host,
            nixl_receiver_port,
            nixl_buffer_size,
            nixl_buffer_device,
            nixl_enable_gc,
            audit_actual_remote_url,
        )

    @staticmethod
    def from_legacy(
        chunk_size: int = 256,
        backend: str = "cpu",
        remote_url: Optional[str] = "lm://localhost:65432",
        remote_serde: str = "naive",
        save_decode_cache: bool = False,
        enable_blending: bool = False,
        blend_recompute_ratio: float = 0.15,
        blend_min_tokens: int = 256,
        blend_special_str: str = " # # ",
        max_local_disk_size: float = 0.0,
        enable_p2p: bool = False,
        lookup_url: Optional[str] = None,
        distributed_url: Optional[str] = None,
        error_handling: bool = False,
        local_disk: str = "",
    ) -> "LMCacheEngineConfig":
        # TODO (ApostaC): Add nixl config
        if backend == "cpu":
            local_cpu = True
            max_local_cpu_size = 5
            local_disk = None
            max_local_disk_size = 0
            remote_url = None
        elif backend == "local_disk":
            local_cpu = False
            max_local_cpu_size = 5
            local_disk = "/local/disk_test/local_disk/"
            max_local_disk_size = 5
            remote_url = None
        elif backend == "local_cpu_disk":
            local_cpu = True
            max_local_cpu_size = 5
            local_disk = "/local/disk_test/local_disk/"
            max_local_disk_size = 5
            remote_url = None
        elif backend == "local_cpu_disk_test":
            local_cpu = True
            max_local_cpu_size = 1
            max_local_disk_size = 1
            remote_url = None
        else:
            raise ValueError(f"Invalid backend: {backend}")

        return LMCacheEngineConfig(
            chunk_size,
            local_cpu,
            max_local_cpu_size,
            local_disk,
            max_local_disk_size,
            remote_url,
            remote_serde,
            save_decode_cache,
            enable_blending,
            blend_recompute_ratio,
            blend_min_tokens,
            blend_special_str,
            enable_p2p,
            lookup_url,
            distributed_url,
            error_handling,
        )
